#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PaddleOCR 批量图片文字识别工具
基于官网示例修改，支持批量处理images文件夹中的图片
OCR模型保持常驻状态，避免重复初始化
"""

import os
import time
from pathlib import Path

from paddleocr import PaddleOCR

# 全局OCR对象，避免重复初始化
ocr = None

def init_ocr():
    """初始化OCR模型（只初始化一次）"""
    global ocr
    if ocr is None:
        print("正在初始化PaddleOCR模型（首次运行需要较长时间）...")
        # ocr = PaddleOCR(
        #     use_doc_orientation_classify=False,
        #     use_doc_unwarping=False,
        #     use_textline_orientation=False
        # )# 文本检测+文本识别
        ocr = PaddleOCR(
            use_doc_orientation_classify=True,
            use_doc_unwarping=True
            )# 文本图像预处理+文本检测+方向分类+文本识别
        # ocr = PaddleOCR(
        #     text_detection_model_name="PP-OCRv5_mobile_det",
        #     text_recognition_model_name="PP-OCRv5_mobile_rec",
        #     use_doc_orientation_classify=False,
        #     use_doc_unwarping=False,
        #     use_textline_orientation=False) # 更换 PP-OCRv5_mobile 模型
        print("PaddleOCR模型初始化完成！\n")
    return ocr

def process_images():
    """处理图片识别"""
    print("=== PaddleOCR 批量图片文字识别工具 ===\n")
    
    # 获取OCR实例
    ocr = init_ocr()
    
    # 设置路径
    input_dir = "images"
    output_dir = "results"
    
    # 创建输出目录
    Path(output_dir).mkdir(exist_ok=True)
    
    # 支持的图片格式
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif', '.webp']
    
    # 获取所有图片文件
    input_path = Path(input_dir)
    if not input_path.exists():
        print(f"错误：图片文件夹不存在: {input_dir}")
        return
    
    image_files = []
    for file_path in input_path.iterdir():
        if file_path.is_file() and file_path.suffix.lower() in image_extensions:
            image_files.append(file_path)
    
    if not image_files:
        print(f"在目录 {input_dir} 中未找到支持的图片文件")
        return
    
    print(f"找到 {len(image_files)} 个图片文件")
    print("开始批量处理...\n")
    
    # 处理每个图片
    for i, image_file in enumerate(image_files, 1):
        print(f"处理进度: {i}/{len(image_files)} - {image_file.name}")
        
        try:
            # 记录开始时间
            start_time = time.time()
            
            # 执行OCR识别
            result = ocr.predict(str(image_file))
            
            # 记录结束时间
            end_time = time.time()
            processing_time = end_time - start_time
            
            # 提取识别结果
            if result and len(result) > 0:
                res = result[0]  # 获取第一个结果
                
                # 从OCRResult对象的json属性中提取数据
                try:
                    # 获取JSON数据
                    json_data = res.json
                    print(f"    JSON数据类型: {type(json_data)}")
                    print(f"    JSON键: {list(json_data.keys())}")
                    
                    # 从res字段中提取数据
                    if 'res' in json_data:
                        res_data = json_data['res']
                        print(f"    res字段类型: {type(res_data)}")
                        print(f"    res字段键: {list(res_data.keys()) if isinstance(res_data, dict) else '不是字典'}")
                        
                        # 提取rec_texts和rec_scores
                        if isinstance(res_data, dict):
                            rec_texts = res_data.get('rec_texts', [])
                            rec_scores = res_data.get('rec_scores', [])
                            print(f"    从res字段找到rec_texts: {len(rec_texts)} 个")
                            print(f"    从res字段找到rec_scores: {len(rec_scores)} 个")
                        else:
                            print(f"    res字段不是字典类型: {type(res_data)}")
                            rec_texts = []
                            rec_scores = []
                    else:
                        print("    未找到res字段")
                        rec_texts = []
                        rec_scores = []
                    
                    print(f"    最终识别到文本: {len(rec_texts)} 个")
                    print(f"    最终置信度分数: {len(rec_scores)} 个")
                    
                    if rec_texts:
                        # 计算平均置信度
                        avg_confidence = sum(rec_scores) / len(rec_scores) if rec_scores else 0.0
                        
                        # 合并所有文本
                        combined_text = '\n'.join(rec_texts)
                        
                        # 生成输出文件名
                        base_name = image_file.stem
                        txt_filename = base_name + '_ocr_result.txt'
                        img_filename = base_name + '_ocr_result.jpg'
                        json_filename = base_name + '_ocr_result.json'
                        
                        # 保存txt文件
                        txt_file_path = Path(output_dir) / txt_filename
                        with open(txt_file_path, 'w', encoding='utf-8') as f:
                            f.write(f"图片文件: {image_file.name}\n")
                            f.write(f"识别时间: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
                            f.write(f"处理耗时: {processing_time:.2f} 秒\n")
                            f.write(f"平均置信度: {avg_confidence:.4f}\n")
                            f.write(f"识别文本数量: {len(rec_texts)}\n")
                            f.write("-" * 50 + "\n")
                            f.write(combined_text)
                        
                        # 保存img和json文件（使用官方方法）
                        try:
                            res.save_to_img(str(output_dir))
                            res.save_to_json(str(output_dir))
                            print(f"  ✓ 识别成功，置信度: {avg_confidence:.4f}，耗时: {processing_time:.2f}秒")
                            print(f"    已保存: {txt_filename}, {img_filename}, {json_filename}")
                        except Exception as save_error:
                            print(f"  ✓ 识别成功，置信度: {avg_confidence:.4f}，耗时: {processing_time:.2f}秒")
                            print(f"    已保存: {txt_filename}，但保存img/json时出错: {save_error}")
                    else:
                        print("  - 未识别到文字")
                        
                except Exception as e:
                    print(f"  ✗ 提取结果失败: {str(e)}")
            else:
                print("  - 未获得识别结果")
                
        except Exception as e:
            print(f"  ✗ 处理失败: {str(e)}")
    
    print(f"\n处理完成！结果已保存到 {output_dir} 文件夹")

def show_menu():
    """显示菜单"""
    print("\n" + "="*50)
    print("PaddleOCR 批量识别工具 - 常驻模式")
    print("="*50)
    print("1. 开始识别图片")
    print("2. 重新初始化模型")
    print("3. 退出程序")
    print("="*50)

def main():
    """主函数 - 交互式菜单"""
    print("正在启动PaddleOCR工具...")
    
    # 首次初始化模型
    init_ocr()
    
    while True:
        show_menu()
        choice = input("请选择操作 (1-3): ").strip()
        
        if choice == '1':
            process_images()
        elif choice == '2':
            global ocr
            ocr = None
            print("模型已重置，下次识别时将重新初始化...")
        elif choice == '3':
            print("感谢使用，再见！")
            break
        else:
            print("无效选择，请重新输入！")

if __name__ == "__main__":
    main()